#include <stdio.h>
#include <vector>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>

#include "net.h"
#if NCNN_VULKAN
#include "gpu.h"
#endif // NCNN_VULKAN

struct Object
{
    cv::Rect_<float> rect;
    int label;
    float prob;
    std::vector<float> maskdata;
    cv::Mat mask;
};

static ncnn::Mat generate_priorbox(int num_priores)
{
    ncnn::Mat priorbox(4, num_priores);
    const int conv_ws[5] = {69, 35, 18, 9, 5};
    const int conv_hs[5] = {69, 35, 18, 9, 5};

    const float aspect_ratios[3] = {1.f, 0.5f, 2.f};
    const float scales[5] = {24.f, 48.f, 96.f, 192.f, 384.f};

    float* pb = priorbox;
    for(int i = 0; i < 5; i++)
    {
        int conv_w = conv_ws[i];
        int conv_h = conv_hs[i];
        int scale = scales[i];
        for(int i = 0; i < conv_h; i++)
        {
            for(int j = 0; j < conv_w; j++)
            {
                float cx = (j + 0.5f) / conv_w;
                float cy = (i + 0.5f) / conv_h;

                for (int k = 0; k < 3; k++)
                {
                    float ar = aspect_ratios[k];

                    ar = sqrt(ar);

                    float w = scale * ar / 550;
                    float h = scale / ar / 550;

                    // The model ncnn provided uses square anchors.
                    // cfg.backbone.use_square_anchors:
                    h = w;

                    pb[0] = cx;
                    pb[1] = cy;
                    pb[2] = w;
                    pb[3] = h;

                    pb += 4;
                }                
            }
        }
    }

    return priorbox;
}

static inline float intersection_area(const Object& a, const Object& b)
{
    cv::Rect_<float> inter = a.rect & b.rect;
    return inter.area();
}


static void fast_nms(std::vector< std::vector<Object> >& class_candidates, std::vector<Object>& objects, const float iou_thresh, const int nms_top_k, const int keep_top_k)
{
    for(int i = 0; i < (int)class_candidates.size(); i++)
    {
        std::vector<Object>& candidate = class_candidates[i];
        //printf("%d %d.\n", i, candidate.size());
        std::sort(candidate.begin(), candidate.end(), 
                    [](const Object& a, const Object& b) {return a.prob > b.prob;});
        if (candidate.size() == 0)
            continue;
        
        if(nms_top_k != 0&& nms_top_k > candidate.size())
        {
            candidate.erase(candidate.begin()+nms_top_k, candidate.end());
        }

        objects.push_back(candidate[0]);
        const int n = candidate.size();
        std::vector<float> areas(n);
        std::vector<int> keep(n);
        for(int j = 0; j < n; j++)
        {
            areas[j] = candidate[j].rect.area();
        }
        std::vector< std::vector<float> > iou_matrix;
        for(int j = 0; j < n; j++)
        {
            std::vector<float> iou_row(n);
            for(int k = 0; k < n; k++)
            {
                float inter_area = intersection_area(candidate[j], candidate[k]);
                float union_area = areas[j] + areas[k] - inter_area;
                iou_row[k] = inter_area / union_area;
            }
            iou_matrix.push_back(iou_row);
        }
        for(int j = 1; j < n; j++)
        {
            std::vector<float>::iterator max_value;
            max_value = std::max_element(iou_matrix[j].begin(), iou_matrix[j].begin()+j-1);
            //printf("%f %d\n", *max_value, max_value-iou_matrix[j].begin());
            if(*max_value <= iou_thresh)
            {
                //printf("%f\n", *max_value);
                objects.push_back(candidate[j]);
            }
        }

    }
    std::sort(objects.begin(), objects.end(), 
                    [](const Object& a, const Object& b) {return a.prob > b.prob;});
    if(objects.size() > keep_top_k)
        objects.resize(keep_top_k);
}

static int detect_yolact(const cv::Mat& bgr, std::vector<Object>& objects)
{
    ncnn::Net yolact;

#if NCNN_VULKAN
    yolact.opt.use_vulkan_compute = true;
#endif // NCNN_VULKAN

    // original model converted from https://github.com/dbolya/yolact
    // yolact_resnet50_54_800000.pth
    // the ncnn model https://github.com/nihui/ncnn-assets/tree/master/models
    yolact.load_param("yolact.param");
    yolact.load_model("yolact.bin");

    const int target_size = 550;

    int img_w = bgr.cols;
    int img_h = bgr.rows;

    ncnn::Mat in = ncnn::Mat::from_pixels_resize(bgr.data, ncnn::Mat::PIXEL_BGR2RGB, img_w, img_h, target_size, target_size);

    const float mean_vals[3] = {123.68f, 116.78f, 103.94f};
    const float norm_vals[3] = {1.0/58.40f, 1.0/57.12f, 1.0/57.38f};
    in.substract_mean_normalize(mean_vals, norm_vals);
    //yolact.dump();
    yolact.input("input.1", in);
    yolact.run();
    ncnn::Mat maskmaps;
    ncnn::Mat location;
    ncnn::Mat mask;
    ncnn::Mat confidence;
    
    yolact.extract("619", maskmaps);// 138x138 x 32

    yolact.extract("816", location);// 4 x 19248
    yolact.extract("818", mask);// maskdim 32 x 19248
    yolact.extract("820", confidence);// 81 x 19248

    int num_class = confidence.w;
    int num_priors = confidence.h;
    ncnn::Mat priorboxes = generate_priorbox(num_priors);
    const float confidence_thresh = 0.05f;
    const float nms_thresh = 0.5f;
    const int keep_top_k = 200;

    std::vector< std::vector<Object> > class_candidates;
    class_candidates.resize(num_class);

    for(int i = 0; i < num_priors; i++)
    {
        const float* conf = confidence.row(i);
        const float* loc = location.row(i);
        const float* maskdata = mask.row(i);
        const float* priorbox = priorboxes.row(i);

        int label = 0;
        float score = 0.f;
        for(int j = 1; j < num_class; j++)
        {
            float class_score = conf[j];
            if(class_score > score)
            {
                label = j;
                score = class_score;
            }
        }

        if(label == 0||score <= confidence_thresh)
            continue;

        float var[4] = {0.1f, 0.1f, 0.2f, 0.2f};

        float bbox_cx = var[0] * loc[0] * priorbox[2] + priorbox[0];
        float bbox_cy = var[1] * loc[1] * priorbox[3] + priorbox[1];
        float bbox_w = (float)(exp(var[2] * loc[2]) * priorbox[2]);
        float bbox_h = (float)(exp(var[3] * loc[3]) * priorbox[3]);
        
        float obj_x1 = bbox_cx - bbox_w * 0.5f;
        float obj_y1 = bbox_cy - bbox_h * 0.5f;
        float obj_x2 = bbox_cx + bbox_w * 0.5f;
        float obj_y2 = bbox_cy + bbox_h * 0.5f;

        obj_x1 = std::max(std::min(obj_x1 * bgr.cols, (float)(bgr.cols - 1)), 0.f);
        obj_y1 = std::max(std::min(obj_y1 * bgr.rows, (float)(bgr.rows - 1)), 0.f);
        obj_x2 = std::max(std::min(obj_x2 * bgr.cols, (float)(bgr.cols - 1)), 0.f);
        obj_y2 = std::max(std::min(obj_y2 * bgr.rows, (float)(bgr.rows - 1)), 0.f);

        Object obj;
        obj.rect = cv::Rect_<float>(obj_x1, obj_y1, obj_x2-obj_x1+1, obj_y2-obj_y1+1);
        obj.label = label;
        obj.prob = score;
        obj.maskdata = std::vector<float>(maskdata, maskdata + mask.w);

        class_candidates[label].push_back(obj);
    }

    objects.clear();
    fast_nms(class_candidates, objects, nms_thresh, 0, keep_top_k);

    for (int i=0; i<objects.size(); i++)
    {
        Object& obj = objects[i];

        cv::Mat mask1(maskmaps.h, maskmaps.w, CV_32FC1);
        {
            mask1 = cv::Scalar(0.f);

            for (int p=0; p<maskmaps.c; p++)
            {
                const float* maskmap = maskmaps.channel(p);
                float coeff = obj.maskdata[p];
                float* mp = (float*)mask1.data;

                // mask += m * coeff
                for (int j=0; j<maskmaps.w * maskmaps.h; j++)
                {
                    mp[j] += maskmap[j] * coeff;
                }
            }
        }

        cv::Mat mask2;
        cv::resize(mask1, mask2, cv::Size(img_w, img_h));

        // crop obj box and binarize
        obj.mask = cv::Mat(img_h, img_w, CV_8UC1);
        {
            obj.mask = cv::Scalar(0);

            for (int y=0; y<img_h; y++)
            {
                if (y < obj.rect.y || y > obj.rect.y + obj.rect.height)
                    continue;

                const float* mp2 = mask2.ptr<const float>(y);
                uchar* bmp = obj.mask.ptr<uchar>(y);

                for (int x=0; x<img_w; x++)
                {
                    if (x < obj.rect.x || x > obj.rect.x + obj.rect.width)
                        continue;

                    bmp[x] = mp2[x] > 0.5f ? 255 : 0;
                }
            }
        }
    }
 
    return 0;
}

static void draw_objects(const cv::Mat& bgr, const std::vector<Object>& objects)
{
    static const char* class_names[] = {"background",
        "person", "bicycle", "car", "motorcycle", "airplane", "bus",
        "train", "truck", "boat", "traffic light", "fire hydrant",
        "stop sign", "parking meter", "bench", "bird", "cat", "dog",
        "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe",
        "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
        "skis", "snowboard", "sports ball", "kite", "baseball bat",
        "baseball glove", "skateboard", "surfboard", "tennis racket",
        "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl",
        "banana", "apple", "sandwich", "orange", "broccoli", "carrot",
        "hot dog", "pizza", "donut", "cake", "chair", "couch",
        "potted plant", "bed", "dining table", "toilet", "tv", "laptop",
        "mouse", "remote", "keyboard", "cell phone", "microwave", "oven",
        "toaster", "sink", "refrigerator", "book", "clock", "vase",
        "scissors", "teddy bear", "hair drier", "toothbrush"};

    static const unsigned char colors[19][3] = {
        {244,  67,  54},
        {233,  30,  99},
        {156,  39, 176},
        {103,  58, 183},
        { 63,  81, 181},
        { 33, 150, 243},
        {  3, 169, 244},
        {  0, 188, 212},
        {  0, 150, 136},
        { 76, 175,  80},
        {139, 195,  74},
        {205, 220,  57},
        {255, 235,  59},
        {255, 193,   7},
        {255, 152,   0},
        {255,  87,  34},
        {121,  85,  72},
        {158, 158, 158},
        { 96, 125, 139}
    };

    cv::Mat image = bgr.clone();

    int color_index = 0;

    for (size_t i = 0; i < objects.size(); i++)
    {
        const Object& obj = objects[i];

        if (obj.prob < 0.15)
            continue;

        fprintf(stderr, "%d = %.5f at %.2f %.2f %.2f x %.2f\n", obj.label, obj.prob,
                obj.rect.x, obj.rect.y, obj.rect.width, obj.rect.height);

        const unsigned char* color = colors[color_index++];

        cv::rectangle(image, obj.rect, cv::Scalar(color[0], color[1], color[2]));

        char text[256];
        sprintf(text, "%s %.1f%%", class_names[obj.label], obj.prob * 100);

        int baseLine = 0;
        cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);

        int x = obj.rect.x;
        int y = obj.rect.y - label_size.height - baseLine;
        if (y < 0)
            y = 0;
        if (x + label_size.width > image.cols)
            x = image.cols - label_size.width;

        cv::rectangle(image, cv::Rect(cv::Point(x, y),
                                      cv::Size(label_size.width, label_size.height + baseLine)),
                      cv::Scalar(255, 255, 255), -1);

        cv::putText(image, text, cv::Point(x, y + label_size.height),
                    cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 0));

        // draw mask
        for (int y=0; y<image.rows; y++)
        {
            const uchar* mp = obj.mask.ptr(y);
            uchar* p = image.ptr(y);
            for (int x=0; x<image.cols; x++)
            {
                if (mp[x] == 255)
                {
                    p[0] = cv::saturate_cast<uchar>(p[0] * 0.5 + color[0] * 0.5);
                    p[1] = cv::saturate_cast<uchar>(p[1] * 0.5 + color[1] * 0.5);
                    p[2] = cv::saturate_cast<uchar>(p[2] * 0.5 + color[2] * 0.5);
                }
                p += 3;
            }
        }
    }

    cv::imwrite("result.png", image);
    cv::imshow("image", image);
    cv::waitKey(0);
}

int main(int argc, char** argv)
{
    if (argc != 2)
    {
        fprintf(stderr, "Usage: %s [imagepath]\n", argv[0]);
        return -1;
    }

    const char* imagepath = argv[1];

    cv::Mat m = cv::imread(imagepath, 1);
    if (m.empty())
    {
        fprintf(stderr, "cv::imread %s failed\n", imagepath);
        return -1;
    }

#if NCNN_VULKAN
    ncnn::create_gpu_instance();
#endif // NCNN_VULKAN

    std::vector<Object> objects;
    detect_yolact(m, objects);

#if NCNN_VULKAN
    ncnn::destroy_gpu_instance();
#endif // NCNN_VULKAN

    draw_objects(m, objects);
    return 0;
}

